import sys
from pathlib import Path

base_path = str(Path(__file__).resolve().parent.parent)
sys.path.append(base_path)
import time
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from create_llm import create_llm

public_txt = "RunnableParallel是LangChain中提高管道运行效率的一个强大工具，合理使用可以显著提高应用程序的性能和响应速度"


def get_title(x):
    llm = create_llm()
    content = llm.invoke(f"从如下内容中提取标题：{public_txt}").content
    return content


def get_key_words(x):
    llm = create_llm()
    content = llm.invoke(f"从如下内容中提取关键词：{public_txt}").content
    return content


parallel_chain = RunnableParallel(
    title=get_title,
    key_words=get_key_words,
    question=lambda x: x.get("input"),
    # question2=RunnablePassthrough(),
    # content=lambda x: len(x["input"]),
)

input_text = "根据标题和关键词生成一篇文章"

temp = parallel_chain.invoke({"input": input_text})
print(temp)


def log(val):
    print(type(val), val)
    return val


prompt = ChatPromptTemplate.from_template(
    "标题：{title},关键字：{key_words},问题：{question}"
)

ai = create_llm()

chain = prompt | ai | StrOutputParser()

full_chain = parallel_chain | log | chain
temp = full_chain.invoke({"input": input_text})
print(temp)
